# !/usr/bin/python3
# -*- coding: utf-8 -*-
# ------------------------------------------
# @Time    : Date - 2021/8/21   Time - 17:31
# @Author  : Spence Guo Tang
# @FileName: data_process.py
# ------------------------------------------

import os
import struct
import torch
import numpy as np
import matplotlib.pyplot as plt


def load_mnist(path, kind='train'):
    """Load MNIST data from `path`"""
    labels_path = os.path.join(path, f'{kind}-labels.idx1-ubyte')
    images_path = os.path.join(path, f'{kind}-images.idx3-ubyte')
    with open(labels_path, 'rb') as lbpath:
        magic, n = struct.unpack('>II', lbpath.read(8))
        labels = np.fromfile(lbpath, dtype=np.uint8)
        labels = labels.astype("int64")
        labels = labels.reshape(len(labels), 1)

    with open(images_path, 'rb') as imgpath:
        magic, num, rows, cols = struct.unpack('>IIII', imgpath.read(16))
        images = np.fromfile(imgpath, dtype=np.uint8).reshape(len(labels), 1, 28, 28)
        images = images.astype("float32")
        # images = images / 255.0

        # 可视化图片示例代码
        # for i in range(5):
        #     plt.matshow(images[i, 0, :, :])
        #     plt.show()

    images = torch.from_numpy(images)/255.0
    labels = torch.from_numpy(labels)
    return images, labels


if __name__ == '__main__':
    images, labels = load_mnist("dataset", kind="t10k")
    # a = np.ones((5, 1, 4, 4))
    # a = a/5
    print()
